import torch
import torch.nn as nn

from monai.metrics import DiceMetric, SurfaceDiceMetric, HausdorffDistanceMetric, MSEMetric, PSNRMetric, MAEMetric, SSIMMetric


class GenMetrics(nn.Module):
    def __init__(self, spatial_dims) -> None:
        super().__init__()
        
        self.mae = MAEMetric()
        self.mse = MSEMetric()
        self.psnr = PSNRMetric(max_val=50)
        self.ssim = SSIMMetric(spatial_dims)
        
    def forward(self, x, y):
        return self.mae(x, y), self.mse(x, y), self.psnr(x, y), self.ssim(x, y)
        

if __name__ == "__main__":
    inps = torch.randn((4, 4, 256, 256))
    out = torch.randn((4, 4, 256, 256))
    
    metric = GenMetrics(2)
    
    for i, j in zip(torch.unbind(inps, 1), torch.unbind(out, 1)):
        print(i.size(), i.unsqueeze(1).size())
        a, b, c, d = metric(i.unsqueeze(1), j.unsqueeze(1))
        print(a.size(), b.size(), c.size(), d.size())
    assert 1 == 2
    
    import numpy as np
    def trans_brats_label(x):
            mask_WT = x.copy()
            mask_WT[mask_WT == 1] = 1
            mask_WT[mask_WT == 2] = 1
            mask_WT[mask_WT == 3] = 1

            mask_TC = x.copy()
            mask_TC[mask_TC == 1] = 1
            mask_TC[mask_TC == 2] = 0
            mask_TC[mask_TC == 3] = 1

            mask_ET = x.copy()
            mask_ET[mask_ET == 1] = 0
            mask_ET[mask_ET == 2] = 0
            mask_ET[mask_ET == 3] = 1
            
            mask = np.stack([mask_WT, mask_TC, mask_ET], axis=0)
            return mask
    
    path = '/Users/qlc/Desktop/Dataset/Brats2023/Adult_Glioma/TrainingData/BraTS-GLI-00127-000/BraTS-GLI-00127-000-seg.nii.gz'
    
    import SimpleITK as sitk
    label = sitk.ReadImage(path)
    label = sitk.GetArrayFromImage(label)
    label = trans_brats_label(label)

    import torch
    label = torch.tensor(label)


    pred = label[:, 35].unsqueeze(0)
    true = label[:, 45].unsqueeze(0)
    
    monai_dice = Metrics(num_classes=3)
    
    out, a, b = monai_dice(pred, true)
    print(out, a, b)
    b = torch.nan_to_num(b, 0.)
    print(b)

    
